#!/usr/bin/env python3
"""
Roboflow Dataset Downloader

This script downloads datasets from Roboflow using their Python SDK.
The datasets will be saved to the current directory.
"""

import os
import sys
import argparse
from roboflow import Roboflow
from config import Config


class RoboflowDownloader:
    """Class to handle Roboflow dataset downloads"""
    
    def __init__(self, api_key=None):
        """Initialize the downloader with API key"""
        self.api_key = api_key or Config.ROBOFLOW_API_KEY
        Config.validate_config()
        self.rf = Roboflow(api_key=self.api_key)
    
    def list_workspaces(self):
        """List all available workspaces"""
        try:
            # The Roboflow SDK doesn't have a direct list_workspaces method
            # This functionality requires using the REST API directly
            print("Note: Workspace listing requires manual specification.")
            print("Please provide your workspace name when downloading datasets.")
            print("You can find your workspace name in your Roboflow dashboard URL.")
            return []
        except Exception as e:
            print(f"Error accessing workspaces: {e}")
            return []
    
    def list_projects(self, workspace_name):
        """List all projects in a workspace"""
        try:
            workspace = self.rf.workspace(workspace_name)
            # The workspace object doesn't have a direct list_projects method
            # This functionality is typically accessed through the web interface
            print(f"Note: Project listing for workspace '{workspace_name}' requires manual specification.")
            print("Please provide your project name when downloading datasets.")
            print("You can find project names in your Roboflow dashboard.")
            return []
        except Exception as e:
            print(f"Error accessing workspace '{workspace_name}': {e}")
            return []
    
    def download_dataset(self, workspace_name, project_name, version=1, 
                        format_type=None, download_dir=None):
        """
        Download a dataset from Roboflow
        
        Args:
            workspace_name (str): Name of the workspace
            project_name (str): Name of the project
            version (int): Version of the dataset (default: 1)
            format_type (str): Format of the dataset (default: yolov8)
            download_dir (str): Directory to save the dataset (default: current directory)
        """
        format_type = format_type or Config.DEFAULT_FORMAT
        download_dir = download_dir or Config.DOWNLOAD_DIR
        
        try:
            print(f"\n=== Starting Dataset Download ===")
            print(f"Workspace: {workspace_name}")
            print(f"Project: {project_name}")
            print(f"Version: {version}")
            print(f"Format: {format_type}")
            print(f"Base directory: {download_dir}")
            
            # Create the expected download folder name
            dataset_folder = f"{project_name}-{version}"
            full_download_path = os.path.join(download_dir, dataset_folder)
            print(f"Dataset will be downloaded to: {full_download_path}")
            
            # Check if dataset already exists
            if os.path.exists(full_download_path):
                existing_contents = os.listdir(full_download_path)
                if any(item in existing_contents for item in ['train', 'valid', 'test', 'data.yaml']):
                    print(f"\n⚠️  Dataset already exists at: {full_download_path}")
                    print(f"Existing contents: {existing_contents}")
                    response = input("Do you want to re-download? (y/n): ")
                    if response.lower() not in ['y', 'yes']:
                        print("Download cancelled.")
                        return True
            
            print(f"\nStep 1: Accessing workspace '{workspace_name}'...")
            workspace = self.rf.workspace(workspace_name)
            print("✓ Workspace accessed successfully")
            
            print(f"\nStep 2: Accessing project '{project_name}'...")
            project = workspace.project(project_name)
            print("✓ Project accessed successfully")
            
            print(f"\nStep 3: Accessing dataset version {version}...")
            dataset = project.version(version)
            print("✓ Dataset version accessed successfully")
            
            # Get dataset info
            if hasattr(dataset, 'id'):
                print(f"Dataset ID: {dataset.id}")
            
            print(f"\nStep 4: Downloading dataset in {format_type} format...")
            print(f"Download location: {full_download_path}")
            print("Please wait while the dataset is being downloaded...\n")
            
            result = dataset.download(format_type, location=full_download_path)
            
            print(f"\n=== Download Completed Successfully! ===")
            print(f"Dataset location: {full_download_path}")
            
            # Verify download
            if os.path.exists(full_download_path):
                contents = os.listdir(full_download_path)
                print(f"Downloaded files: {contents}")
                
                # Check for typical YOLO structure
                if 'data.yaml' in contents:
                    data_yaml_path = os.path.join(full_download_path, 'data.yaml')
                    print(f"Configuration file: {data_yaml_path}")
                
                # Count images in train folder if exists
                train_path = os.path.join(full_download_path, 'train', 'images')
                if os.path.exists(train_path):
                    train_images = len([f for f in os.listdir(train_path) if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
                    print(f"Training images: {train_images}")
            
            return True
            
        except Exception as e:
            print(f"\n✗ Error downloading dataset: {e}")
            print(f"Error type: {type(e).__name__}")
            return False
    
    def get_dataset_info(self, workspace_name, project_name, version=1):
        """Get information about a dataset"""
        try:
            print(f"\n=== Getting Dataset Information ===")
            print(f"Workspace: {workspace_name}")
            print(f"Project: {project_name}")
            print(f"Version: {version}")
            
            workspace = self.rf.workspace(workspace_name)
            project = workspace.project(project_name)
            dataset = project.version(version)
            
            print(f"\n=== Dataset Details ===")
            if hasattr(dataset, 'id'):
                print(f"Dataset ID: {dataset.id}")
            
            # Check if dataset already exists locally
            dataset_folder = f"{project_name}-{version}"
            local_path = os.path.join(Config.DOWNLOAD_DIR, dataset_folder)
            
            if os.path.exists(local_path):
                print(f"\n✓ Dataset found locally at: {local_path}")
                contents = os.listdir(local_path)
                print(f"Contents: {contents}")
                
                # Check data.yaml for class information
                data_yaml_path = os.path.join(local_path, 'data.yaml')
                if os.path.exists(data_yaml_path):
                    try:
                        import yaml
                        with open(data_yaml_path, 'r') as f:
                            data_config = yaml.safe_load(f)
                        if 'names' in data_config:
                            print(f"Classes ({data_config.get('nc', 'unknown')}): {data_config['names']}")
                    except:
                        print("Could not read class information from data.yaml")
            else:
                print(f"\n⚠️  Dataset not found locally. Expected location: {local_path}")
                print("Use download command to get the dataset.")
            
            return dataset
            
        except Exception as e:
            print(f"\n✗ Error getting dataset info: {e}")
            return None


def main():
    """Main function to handle command line arguments and execute downloads"""
    parser = argparse.ArgumentParser(description='Download datasets from Roboflow')
    
    parser.add_argument('--list-workspaces', action='store_true',
                       help='List all available workspaces')
    
    parser.add_argument('--list-projects', type=str, metavar='WORKSPACE',
                       help='List all projects in the specified workspace')
    
    parser.add_argument('--workspace', '-w', type=str, required=False,
                       help='Workspace name')
    
    parser.add_argument('--project', '-p', type=str, required=False,
                       help='Project name')
    
    parser.add_argument('--version', '-v', type=int, default=1,
                       help='Dataset version (default: 1)')
    
    parser.add_argument('--format', '-f', type=str, default=Config.DEFAULT_FORMAT,
                       help=f'Dataset format (default: {Config.DEFAULT_FORMAT})')
    
    parser.add_argument('--output', '-o', type=str, default=Config.DOWNLOAD_DIR,
                       help='Output directory (default: current directory)')
    
    parser.add_argument('--info', action='store_true',
                       help='Show dataset information without downloading')
    
    args = parser.parse_args()
    
    # Validate API key
    try:
        Config.validate_config()
    except ValueError as e:
        print(f"Configuration error: {e}")
        print("Please check your .env file or set ROBOFLOW_API_KEY environment variable.")
        sys.exit(1)
    
    downloader = RoboflowDownloader()
    
    # Handle different commands
    if args.list_workspaces:
        downloader.list_workspaces()
        return
    
    if args.list_projects:
        downloader.list_projects(args.list_projects)
        return
    
    if args.workspace and args.project:
        if args.info:
            downloader.get_dataset_info(args.workspace, args.project, args.version)
        else:
            success = downloader.download_dataset(
                workspace_name=args.workspace,
                project_name=args.project,
                version=args.version,
                format_type=args.format,
                download_dir=args.output
            )
            if not success:
                sys.exit(1)
    else:
        print("Error: Both --workspace and --project are required for downloading.")
        print("Use --help for more information.")
        sys.exit(1)


if __name__ == "__main__":
    main()